{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 目标跟踪 - 演示示例\n",
    "\n",
    "在本示例中，我们将展示如何使用JetBot跟随对象！我们将使用预训练的神经网络\n",
    "在[COCO数据集]（http://cocodataset.org）上进行了训练，以检测90个不同的常见对象。这些包括\n",
    "\n",
    "*人（索引0）\n",
    "*杯（索引47）\n",
    "\n",
    "和许多其他文件（您可以检查[此文件]（https://github.com/tensorflow/models/blob/master/research/object_detection/data/mscoco_complete_label_map.pbtxt）以获取类索引的完整列表）。该模型源自[TensorFlow对象检测API]（https://github.com/tensorflow/models/tree/master/research/object_detection），\n",
    "该工具还提供了用于训练自定义任务的对象检测器的实用程序！训练完模型后，我们将在Jetson Nano上使用NVIDIA TensorRT对其进行优化。\n",
    "\n",
    "这使网络非常快速，能够在Jetson Nano上实时执行！但是，我们不会在此笔记本中完成所有的培训和优化步骤。\n",
    "\n",
    "无论如何，让我们开始吧。首先，我们要导入采用我们预先训练的SSD引擎的ObjectDetector类。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 加载物体检测模型\n",
    "> 此处我们提供ssd_mobilenet_v2_coco模型的百度云下载链接：https://pan.baidu.com/s/1fzkK2QWgZqtw85HFUVDAqw  提取码：493j"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "outputs": [],
   "source": [
    "from jetbot import ObjectDetector\n",
    "\n",
    "model = ObjectDetector('ssd_mobilenet_v2_coco.engine')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "ObjectDetector类内部使用TensorRT Python API执行我们提供的引擎。 它还负责对神经网络的输入进行预处理，因为\n",
    "以及解析检测到的对象。 目前，它仅适用于使用jetbot.ssd_tensorrt软件包创建的引擎。 该软件包具有用于转换的实用程序\n",
    "从TensorFlow对象检测API到优化的TensorRT引擎的模型。\n",
    "\n",
    "接下来，让我们初始化相机。 我们的检测器需要300x300像素输入，因此我们在创建相机时会进行设置。\n",
    "\n",
    ">在内部，Camera类使用GStreamer来利用Jetson Nano的图像信号处理器（ISP）。 这是超快的，卸载\n",
    ">许多来自CPU的调整大小计算。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from jetbot import Camera\n",
    "\n",
    "camera = Camera.instance(width=300, height=300)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "现在，让我们使用一些摄像机输入来执行我们的网络。 默认情况下，“ ObjectDetector”类期望相机生成的“ bgr8”格式。 然而，\n",
    "如果您输入的格式不同，则可以覆盖默认的预处理功能。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "detections = model(camera.value)\n",
    "\n",
    "print(detections)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "如果相机视野中有任何COCO对象，则应将其存储在``detections''变量中。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 在文本区域中显示检测结果\n",
    "\n",
    "我们将使用下面的代码打印出检测到的对象。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import display\n",
    "import ipywidgets.widgets as widgets\n",
    "\n",
    "detections_widget = widgets.Textarea()\n",
    "\n",
    "detections_widget.value = str(detections)\n",
    "\n",
    "display(detections_widget)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "您应该看到每个图像中检测到的每个对象的标签，置信度和边界框。 在此示例中，只有一幅图像（我们的相机）。\n",
    "\n",
    "\n",
    "要仅打印在第一张图像中检测到的第一个对象，我们可以调用以下命令\n",
    "\n",
    ">如果未检测到任何对象，则可能会引发错误"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "image_number = 0\n",
    "object_number = 0\n",
    "\n",
    "print(detections[image_number][object_number])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 控制机器人跟随中心物体\n",
    "\n",
    "现在，我们希望机器人跟随指定类的对象。 为此，我们将执行以下操作\n",
    "\n",
    "1.检测与指定类别匹配的对象\n",
    "2.选择最接近摄像机视场中心的对象，这是“目标”对象\n",
    "3.将机器人转向目标对象，否则会漂移\n",
    "4.如果我们被障碍物挡住，请向左转\n",
    "\n",
    "我们还将创建一些小部件，这些小部件将用于控制目标对象标签，机械手速度和\n",
    "“转弯增益”，它将根据目标物体之间的距离控制机器人转弯的速度\n",
    "以及机器人视场的中心。\n",
    "\n",
    "\n",
    "首先，让我们加载碰撞检测模型。 为了方便起见，预先训练的模型存储在此目录中，但是如果您遵循\n",
    "如果针对机器人环境进行了更好的调整，则可能需要使用该模型来避免碰撞。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torchvision\n",
    "import torch.nn.functional as F\n",
    "import cv2\n",
    "import numpy as np\n",
    "\n",
    "collision_model = torchvision.models.alexnet(pretrained=False)\n",
    "collision_model.classifier[6] = torch.nn.Linear(collision_model.classifier[6].in_features, 2)\n",
    "collision_model.load_state_dict(torch.load('../collision_avoidance/best_model.pth'))\n",
    "device = torch.device('cuda')\n",
    "collision_model = collision_model.to(device)\n",
    "\n",
    "mean = 255.0 * np.array([0.485, 0.456, 0.406])\n",
    "stdev = 255.0 * np.array([0.229, 0.224, 0.225])\n",
    "\n",
    "normalize = torchvision.transforms.Normalize(mean, stdev)\n",
    "\n",
    "def preprocess(camera_value):\n",
    "    global device, normalize\n",
    "    x = camera_value\n",
    "    x = cv2.resize(x, (224, 224))\n",
    "    x = cv2.cvtColor(x, cv2.COLOR_BGR2RGB)\n",
    "    x = x.transpose((2, 0, 1))\n",
    "    x = torch.from_numpy(x).float()\n",
    "    x = normalize(x)\n",
    "    x = x.to(device)\n",
    "    x = x[None, ...]\n",
    "    return x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "太好了，现在让我们初始化我们的机器人，以便我们可以控制电机。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from jetbot import Robot\n",
    "\n",
    "robot = Robot()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "最后，让我们显示所有控件小部件，并将网络执行功能连接到摄像机更新。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from jetbot import bgr8_to_jpeg\n",
    "\n",
    "blocked_widget = widgets.FloatSlider(min=0.0, max=1.0, value=0.0, description='blocked')\n",
    "image_widget = widgets.Image(format='jpeg', width=300, height=300)\n",
    "label_widget = widgets.IntText(value=1, description='tracked label')\n",
    "speed_widget = widgets.FloatSlider(value=0.4, min=0.0, max=1.0, description='speed')\n",
    "turn_gain_widget = widgets.FloatSlider(value=0.8, min=0.0, max=2.0, description='turn gain')\n",
    "\n",
    "display(widgets.VBox([\n",
    "    widgets.HBox([image_widget, blocked_widget]),\n",
    "    label_widget,\n",
    "    speed_widget,\n",
    "    turn_gain_widget\n",
    "]))\n",
    "\n",
    "width = int(image_widget.width)\n",
    "height = int(image_widget.height)\n",
    "\n",
    "def detection_center(detection):\n",
    "    \"\"\"Computes the center x, y coordinates of the object\"\"\"\n",
    "    bbox = detection['bbox']\n",
    "    center_x = (bbox[0] + bbox[2]) / 2.0 - 0.5\n",
    "    center_y = (bbox[1] + bbox[3]) / 2.0 - 0.5\n",
    "    return (center_x, center_y)\n",
    "    \n",
    "def norm(vec):\n",
    "    \"\"\"Computes the length of the 2D vector\"\"\"\n",
    "    return np.sqrt(vec[0]**2 + vec[1]**2)\n",
    "\n",
    "def closest_detection(detections):\n",
    "    \"\"\"Finds the detection closest to the image center\"\"\"\n",
    "    closest_detection = None\n",
    "    for det in detections:\n",
    "        center = detection_center(det)\n",
    "        if closest_detection is None:\n",
    "            closest_detection = det\n",
    "        elif norm(detection_center(det)) < norm(detection_center(closest_detection)):\n",
    "            closest_detection = det\n",
    "    return closest_detection\n",
    "        \n",
    "def execute(change):\n",
    "    image = change['new']\n",
    "    \n",
    "    # execute collision model to determine if blocked\n",
    "    collision_output = collision_model(preprocess(image)).detach().cpu()\n",
    "    prob_blocked = float(F.softmax(collision_output.flatten(), dim=0)[0])\n",
    "    blocked_widget.value = prob_blocked\n",
    "    \n",
    "    # turn left if blocked\n",
    "    if prob_blocked > 0.5:\n",
    "        robot.left(0.3)\n",
    "        image_widget.value = bgr8_to_jpeg(image)\n",
    "        return\n",
    "        \n",
    "    # compute all detected objects\n",
    "    detections = model(image)\n",
    "    \n",
    "    # draw all detections on image\n",
    "    for det in detections[0]:\n",
    "        bbox = det['bbox']\n",
    "        cv2.rectangle(image, (int(width * bbox[0]), int(height * bbox[1])), (int(width * bbox[2]), int(height * bbox[3])), (255, 0, 0), 2)\n",
    "    \n",
    "    # select detections that match selected class label\n",
    "    matching_detections = [d for d in detections[0] if d['label'] == int(label_widget.value)]\n",
    "    \n",
    "    # get detection closest to center of field of view and draw it\n",
    "    det = closest_detection(matching_detections)\n",
    "    if det is not None:\n",
    "        bbox = det['bbox']\n",
    "        cv2.rectangle(image, (int(width * bbox[0]), int(height * bbox[1])), (int(width * bbox[2]), int(height * bbox[3])), (0, 255, 0), 5)\n",
    "    \n",
    "    \n",
    "        \n",
    "    # otherwise go forward if no target detected\n",
    "    if det is None:\n",
    "        robot.forward(float(speed_widget.value))\n",
    "        \n",
    "    # otherwsie steer towards target\n",
    "    else:\n",
    "        # move robot forward and steer proportional target's x-distance from center\n",
    "        center = detection_center(det)\n",
    "        robot.set_motors(\n",
    "            float(speed_widget.value + turn_gain_widget.value * center[0]),\n",
    "            float(speed_widget.value - turn_gain_widget.value * center[0])\n",
    "        )\n",
    "    \n",
    "    # update image widget\n",
    "    image_widget.value = bgr8_to_jpeg(image)\n",
    "    \n",
    "execute({'new': camera.value})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "调用下面的块，将执行功能连接到每个相机框架更新。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "camera.unobserve_all()\n",
    "camera.observe(execute, names='value')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "太棒了！ 如果未阻止机器人，则应该在检测到的对象周围看到蓝色框。 目标对象（机器人遵循的对象）将显示为绿色。\n",
    "\n",
    "被检测到时，机器人应转向目标。 如果它被物体挡住，它只会向左转。\n",
    "\n",
    "您可以调用下面的代码块来手动断开相机的处理并停止机器人。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "\n",
    "camera.unobserve_all()\n",
    "time.sleep(1.0)\n",
    "robot.stop()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
